from mlir.ir import *


class Toy:
    WeightOp = 'toy.Weight'
    ConstOp = 'toy.const'
    InputOp = 'toy.input'
    AddOp = 'toy.add'
    ConvOp = 'toy.conv2d'
    MatMulOp = 'toy.matmul'
    MaxPoolOp = 'toy.max_pool2d'
    ReshapeOp = 'toy.reshape'
    ReluOp = 'toy.relu'
    PrintOp = 'toy.print'
    FuncOp = 'toy.func'

class State:
    TOP_F32 = 'TOP_F32'


def get_weight_file(model_name: str, state: str):
    name = "{}_{}_weight.npz".format(model_name, state)
    return name.lower()


def checkType(obj, type):
    if not isinstance(obj, type):
        raise AttributeError('{} is not {}'.format(obj, type))


class MLIRImporter(object):

    def __init__(self,
                 input_shapes: list,
                 output_shapes: list,
                 model_name: str,
                 input_types: list = [],
                 output_types: list = [],
                 state: str = State.TOP_F32,
                 do_declare: bool = True):
        """
            input_shape: List[List], put module input shape. ex: [[1, 3, 224, 224]]
            output_shape: List, put module output shape. ex: [1, 1000]
        """
        assert (len(model_name) > 0)
        self.model_name = model_name
        self.state = state
        self.ctx = Context()
        self.ctx.allow_unregistered_dialects = True
        self.loc = Location.unknown(self.ctx)
        self.ctx.__enter__()
        self.loc.__enter__()
        self.input_shapes = input_shapes
        self.output_shapes = output_shapes
        self.num_input = len(self.input_shapes)
        self.num_output = len(self.output_shapes)
        self.load_weight = dict()
        self.mlir_type = {
            "INT8": IntegerType.get_signless(8),
            "UINT8": IntegerType.get_unsigned(8),
            "SINT8": IntegerType.get_signed(8),
            "INT16": IntegerType.get_signless(16),
            "INT32": IntegerType.get_signless(32),
            "INT64": IntegerType.get_signless(64),
            "F64": F64Type.get(),
            "F32": F32Type.get(),
            "F16": F16Type.get(),
            "BF16": BF16Type.get()
        }
        if do_declare:
            self.declare_func(input_types, output_types)

    def __del__(self):
        try:
            self.loc.__exit__(None, None, None)
        except:
            pass
        try:
            self.ctx.__exit__(None, None, None)
        except:
            pass

    def ArrayAttr(self, data: list, data_type: str = 'INT64'):
        assert (data_type in self.mlir_type)
        if data_type.find("INT") >= 0:
            return ArrayAttr.get([IntegerAttr.get(self.mlir_type[data_type], x) for x in data])
        if data_type == 'F32':
            return ArrayAttr.get([FloatAttr.get_f32(x) for x in data])
        if data_type == 'F64':
            return ArrayAttr.get([FloatAttr.get_f64(x) for x in data])
        raise RuntimeError("unsupport data type:{}".format(data_type))

    def get_value_type(self, value):
        _type = str(value.type)
        _type = _type.split('<')[-1].split('x')[-1].split('>')[0]
        if _type == "f32":
            return self.mlir_type['F32']
        if _type == "f64":
            return self.mlir_type['F64']
        elif _type == "i8":
            return self.mlir_type['INT8']
        elif _type == "ui8":
            return self.mlir_type['UINT8']
        else:
            raise RuntimeError("No support {}".format(_type))

    def buildOp(self, op_type, operands, output_types: list, **kargs):
        """
            op_type: String
            inputOpreands: List[pybind.op]
            output_types: List[pybind.op]
            kargs: Dict
        """
        num_output = len(output_types)
        if num_output == 1:
            loc = Location.fused([Location.name(kargs['name'])])
        else:
            assert (isinstance(kargs["name"], list))
            assert (num_output == len(kargs["name"]))
            loc = Location.fused([Location.name(n) for n in kargs["name"]])
        del kargs["name"]
        op = Operation.create(
            op_type,
            results=output_types,
            operands=operands,
            loc=loc,
            attributes=kargs,
        )
        self.insert_point.insert(op)

        assert (num_output == len(output_types))
        if num_output > 1:
            return tuple(op.results)
        else:
            return op.result

    def create_input_op(self, name, index, **kargs):
        assert (index < len(self.func_args))
        param = {}
        if 'channel_format' in kargs:
            param['channel_format'] = StringAttr.get(kargs['channel_format'])
        op = Operation.create(Toy.InputOp,
                              results=[self.input_types[index]],
                              operands=[self.func_args[index]],
                              loc=Location.fused([Location.name(name)]),
                              attributes=param)
        self.insert_point.insert(op)
        return op.results[0]

    def create_weight_op(self, name, weight_data, output_shape, data_type="F64"):
        if name in self.load_weight:
            _op, _shape, _type = self.load_weight[name]
            if _shape != output_shape or _type != data_type:
                raise RuntimeError("{} weight conflict".format(name))
            return _op
        tensor_type = RankedTensorType.get(output_shape, self.mlir_type[data_type])
        f64_weight_data = weight_data.astype(np.float64)
        tensor_data = DenseElementsAttr.get(f64_weight_data, type = F64Type.get(), context = self.ctx)
        op = Operation.create(Toy.ConstOp,
                              results=[tensor_type],
                              attributes = {'value': tensor_data},
                              loc=Location.fused([Location.name(name)]))

        self.insert_point.insert(op)
        result = op.results[0]
        self.load_weight[name] = (result, output_shape, data_type)
        return result

    def create_matmul_op(self, operands, output_shape, **kargs):
        # get_value_type
        output_type = RankedTensorType.get(tuple(output_shape), self.get_value_type(operands[0]))

        param = {
            'name': kargs['name'],
            'do_relu': BoolAttr.get(kargs['do_relu']),
        }
        return self.buildOp(Toy.MatMulOp, operands, [output_type], **param)

    def create_relu_op(self, operands, output_shape, **kargs):
        output_type = RankedTensorType.get(tuple(output_shape), self.get_value_type(operands[0]))
        param = {
            'name': kargs['name'],
        }
        if 'relu_limit' in kargs:
            param['relu_limit'] = FloatAttr.get_f64(kargs['relu_limit'])
        return self.buildOp(Toy.ReluOp, operands, [output_type], **param)

    def create_return_op(self, Operands):
        return_op = Operation.create("func.return", operands=Operands, results=[])
        self.insert_point.insert(return_op)
        return return_op

    def create_print_op(self, Operands):
        print_op = Operation.create(Toy.PrintOp, operands= Operands, results=[])
        self.insert_point.insert(print_op)
        return print_op

    def create_reshape_op(self, operands, output_shape, **kargs):
        output_type = RankedTensorType.get(tuple(output_shape), self.get_value_type(operands[0]))
        return self.buildOp(Toy.ReshapeOp, operands, [output_type], name=kargs['name'])

    def print_module(self):
        mlir_format = self.mlir_module.operation.get_asm(enable_debug_info=True)
        return mlir_format

    def declare_func(self, input_types: list = [], output_types: list = []):
        if len(input_types) == 0:
            input_types = self.num_input * ['F64']
        if len(output_types) == 0:
            output_types = self.num_output * ['F64']
        self.input_types = list()
        self.output_types = list()
        for _shape, _type in zip(self.input_shapes, input_types):
            if isinstance(_type, str):
                self.input_types.append(RankedTensorType.get(_shape, self.mlir_type[_type]))
            else:
                self.input_types.append(RankedTensorType.get(_shape, _type))
        for _shape, _type in zip(self.output_shapes, output_types):
            if isinstance(_type, str):
                self.output_types.append(RankedTensorType.get(_shape, self.mlir_type[_type]))
            else:
                self.output_types.append(RankedTensorType.get(_shape, _type))
        args_txt = str()
        for _idx, _type in enumerate(self.input_types):
            args_txt += "%args{}: {} loc(unknown)".format(_idx, _type.__str__())
            if (_idx + 1) < self.num_input:
                args_txt += ", "
        output_txt = str()
        for _idx, _type in enumerate(self.output_types):
            output_txt += _type.__str__()
            if (_idx + 1) < self.num_output:
                output_txt += ", "
        if self.num_output > 1:
            output_txt = "({})".format(output_txt)

        main_func = """
            module attributes {{module.name = \"{name}\"}} {{
                func.func @main({args}) {{
                    %0 = \"toy.none\"() : () -> none
                }}
            }}
        """.format(name=self.model_name, args=args_txt)

        self.mlir_module = Module.parse(main_func, self.ctx)
        
        self.mlir_module.dump()
        self.func = self.mlir_module.body.operations[0]
        self.entry_block = self.func.regions[0].blocks[0]
        self.insert_point = InsertionPoint(self.entry_block)
        self.none_op = self.entry_block.operations[0].operation.results[0]

        self.func_args = list()
        for i in self.entry_block.arguments:
            self.func_args.append(i)
